from python_a2a.server import A2AServer
from python_a2a.mcp import FastMCPAgent
from python_a2a.models import Message, TextContent, MessageRole
from python_a2a import run_server
import re
import asyncio
import logging

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# DuckDuckGo Agent for ticker lookup
class DuckDuckGoAgent(A2AServer, FastMCPAgent):
    """Agent that finds stock ticker symbols."""
    
    def __init__(self):
        super().__init__(url="http://localhost:5003")
        FastMCPAgent.__init__(
            self,
            mcp_servers={"search": "http://localhost:5001"}
        )
    
    def handle_message(self, message):
        """同步处理消息"""
        logger.info(f"处理DuckDuckGo消息: {message.content.text}")
        
        try:
            if hasattr(message.content, 'text'):
                text = message.content.text
                
                # 提取公司名称
                company_name = self._extract_company_name(text)
                
                if not company_name:
                    return Message(
                        content=TextContent(text="请提供公司名称来查找股票代码。"),
                        role=MessageRole.AGENT,
                        parent_message_id=getattr(message, 'message_id', None),
                        conversation_id=getattr(message, 'conversation_id', None)
                    )
                
                # 同步调用MCP工具查找股票代码
                try:
                    # 由于MCP工具可能是异步的，我们需要在同步环境中运行异步代码
                    ticker = asyncio.run(self._search_ticker_async(company_name))
                    
                    return Message(
                        content=TextContent(text=f"{company_name} 的股票代码是 {ticker}"),
                        role=MessageRole.AGENT,
                        parent_message_id=getattr(message, 'message_id', None),
                        conversation_id=getattr(message, 'conversation_id', None)
                    )
                    
                except Exception as e:
                    logger.error(f"调用MCP工具时出错: {e}")
                    # 返回模拟数据作为备用
                    simulated_ticker = self._simulate_ticker_lookup(company_name)
                    return Message(
                        content=TextContent(text=f"{company_name} 的股票代码是 {simulated_ticker} (模拟数据)"),
                        role=MessageRole.AGENT,
                        parent_message_id=getattr(message, 'message_id', None),
                        conversation_id=getattr(message, 'conversation_id', None)
                    )
        
        except Exception as e:
            logger.error(f"处理消息时出错: {e}")
            return Message(
                content=TextContent(text="处理您的请求时出现了错误。"),
                role=MessageRole.AGENT,
                parent_message_id=getattr(message, 'message_id', None),
                conversation_id=getattr(message, 'conversation_id', None)
            )
        
        # 默认响应
        return Message(
            content=TextContent(text="我可以帮助查找公司的股票代码。"),
            role=MessageRole.AGENT,
            parent_message_id=getattr(message, 'message_id', None),
            conversation_id=getattr(message, 'conversation_id', None)
        )
    
    async def _search_ticker_async(self, company_name: str):
        """异步搜索股票代码"""
        try:
            # 调用MCP工具查找股票代码
            ticker = await self.call_mcp_tool("search", "search_ticker", company_name=company_name)
            return ticker
        except Exception as e:
            logger.error(f"MCP工具调用失败: {e}")
            # 返回模拟数据
            return self._simulate_ticker_lookup(company_name)
    
    def _extract_company_name(self, text: str) -> str:
        """从文本中提取公司名称"""
        # 多种匹配模式
        patterns = [
            r"(.+?)的股票代码",
            r"股票代码(.+?)",
            r"(.+?)的股票",
            r"(.+?)股票代码",
            r"ticker\s+(?:for|of)\s+(.+)",
            r"stock\s+(?:symbol|ticker)\s+(?:for|of)\s+(.+)",
            r"symbol\s+for\s+(.+)"
        ]
        
        text_lower = text.lower()
        
        for pattern in patterns:
            match = re.search(pattern, text_lower, re.IGNORECASE)
            if match:
                company = match.group(1).strip()
                # 清理结果
                company = re.sub(r'[?？。.,!！]', '', company)
                return company
        
        # 如果没有匹配模式，返回整个文本（去除常见疑问词）
        cleaned_text = re.sub(r'(什么|多少|哪个|哪只|查找|查询|搜索|的股票代码|的股票|股票代码|ticker|stock|symbol)', '', text_lower, flags=re.IGNORECASE)
        cleaned_text = cleaned_text.strip()
        
        return cleaned_text if cleaned_text else ""
    
    def _simulate_ticker_lookup(self, company_name: str) -> str:
        """模拟股票代码查找（备用方案）"""
        ticker_mapping = {
            "apple": "AAPL",
            "苹果": "AAPL",
            "microsoft": "MSFT", 
            "微软": "MSFT",
            "google": "GOOGL",
            "谷歌": "GOOGL",
            "amazon": "AMZN",
            "亚马逊": "AMZN",
            "tesla": "TSLA",
            "特斯拉": "TSLA",
            "meta": "META",
            "facebook": "META",
            "脸书": "META",
            "阿里巴巴": "BABA",
            "腾讯": "TCEHY",
            "百度": "BIDU",
            "茅台": "600519",
            "美团": "3690.HK"
        }
        
        # 查找匹配的公司
        for key, value in ticker_mapping.items():
            if key in company_name.lower():
                return value
        
        # 如果没有匹配，返回一个模拟的股票代码
        return f"{company_name.upper()[:4]}"

if __name__ == "__main__":
    # 在5003端口运行此代理
    logger.info("DuckDuckGo股票代码查找代理正在 http://localhost:5003 上运行")
    run_server(DuckDuckGoAgent(), port=5003, debug=True)